# model_ChineseCLIPSimilarityCalculator.py
"""
使用 OFA-Sys/chinese-clip-large 模型计算图像和文本相似度的计算器。
需要 transformers 库。
此文件用于替换原始的、基于 cn-clip 库的实现。
"""

import torch
# Import modules specifically for the OFA-Sys Chinese CLIP model
from transformers import ChineseCLIPModel, ChineseCLIPProcessor
import torch.nn.functional as F
from PIL import Image
import os
import numpy as np


class ChineseCLIPSimilarityCalculator:
    """
    使用 OFA-Sys/chinese-clip-large 模型计算图像和文本相似度的计算器。
    需要 transformers 库。
    """
    def __init__(self, clip_model_path, device="cuda"):
        """
        初始化计算器。
        Args:
            clip_model_path (str): 本地 chinese-clip-large 模型的路径。
            device (str): 运行模型的设备 ('cuda' or 'cpu')。
        """
        self.device = device
        print("------------")
        print(f"[ChineseCLIPSimilarityCalculator] Loading processor from {clip_model_path}")
        # 使用 ChineseCLIPProcessor 加载处理器 (用于图像和文本预处理)
        self.processor = ChineseCLIPProcessor.from_pretrained(clip_model_path)
        print(f"[ChineseCLIPSimilarityCalculator] Loading model from {clip_model_path}")
        # 使用 ChineseCLIPModel 加载模型
        self.model = ChineseCLIPModel.from_pretrained(clip_model_path).to(self.device)
        self.model.eval()
        print("[ChineseCLIPSimilarityCalculator] Model loaded and set to eval mode.")

    def _prepare_image(self, image):
        """将输入转换为 PIL Image 对象"""
        if isinstance(image, Image.Image):
            return image
        elif isinstance(image, str) and os.path.exists(image):
            return Image.open(image).convert("RGB")
        elif isinstance(image, torch.Tensor):
            if image.dim() == 4:
                image = image[0].cpu()
            # 确保 Tensor 值在 [0, 1] 范围内再转换
            if image.max() > 1.0:
                 image = image / 255.0
            image_np = (image.permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8)
            return Image.fromarray(image_np)
        elif isinstance(image, np.ndarray):
            if image.dtype != np.uint8:
                # 假设输入是 float [0, 1] 或 [0, 255]
                if image.max() > 1.0:
                    image = image / 255.0
                image = (image * 255).astype(np.uint8)
            if image.ndim == 3 and image.shape[-1] == 1:
                image = np.repeat(image, 3, axis=-1) # 单通道转RGB
            elif image.ndim == 2: # 灰度图
                 image = np.stack([image]*3, axis=-1)
            return Image.fromarray(image).convert("RGB")
        else:
            raise TypeError(f"[ChineseCLIPSimilarityCalculator] Unsupported image type: {type(image)}")

    def get_image_embedding(self, image):
        """
        获取图像的嵌入向量。
        Args:
            image: PIL Image, 图像路径(str), PyTorch Tensor, 或 NumPy array。
        Returns:
            torch.Tensor: 归一化的图像特征向量 [1, D]。
        """
        with torch.no_grad(): # Disable autograd and potentially amp for consistency
            pil_image = self._prepare_image(image)
            # 使用 processor 处理图像
            inputs = self.processor(images=pil_image, return_tensors="pt").to(self.device)
            # 获取图像特征
            image_features = self.model.get_image_features(**inputs)
            # 返回 L2 归一化后的特征
            return F.normalize(image_features, dim=-1)

    def get_text_embeddings(self, texts):
        """
        获取文本的嵌入向量。
        Args:
            texts (str or list[str]): 单个文本字符串或文本列表。
        Returns:
            torch.Tensor: 归一化的文本特征向量 [N, D]。
        """
        with torch.no_grad():
            if isinstance(texts, str):
                texts = [texts]
            # 使用 processor 处理文本
            inputs = self.processor(text=texts, padding=True, truncation=True, return_tensors="pt").to(self.device)
            # 获取文本特征
            text_features = self.model.get_text_features(**inputs)
            # 返回 L2 归一化后的特征
            return F.normalize(text_features, dim=-1)

    def calculate_similarity(self, image, texts):
        """
        计算图像与文本列表的余弦相似度。
        Args:
            image: PIL Image, 图像路径(str), PyTorch Tensor, 或 NumPy array。
            texts (list[str]): 文本列表。
        Returns:
            np.ndarray: 相似度数组 [len(texts)]。
        """
        if not texts: # Handle empty input gracefully
             return np.array([])

        with torch.no_grad():
            image_features = self.get_image_embedding(image) # [1, D]
            text_features = self.get_text_embeddings(texts)  # [N, D]

            # 计算余弦相似度 (点积，因为特征已归一化)
            # image_features: [1, D], text_features.T: [D, N] -> similarities: [1, N]
            similarities = torch.matmul(image_features, text_features.T)
            # --- 关键修改：转换为 float32 再转 numpy ---
            # 转换为 numpy 并展平为 [N]
            # return similarities.cpu().numpy().flatten() # 旧的，可能导致错误
            return similarities.cpu().to(torch.float32).numpy().flatten() # 新的，避免 BF16 错误


    def filter_texts_by_similarity(self, image, texts, threshold=0.3, verbose=False):
        """
        根据相似度阈值过滤文本。
        Args:
            image: PIL Image, 图像路径(str), PyTorch Tensor, 或 NumPy array。
            texts (list[str]): 文本列表。
            threshold (float): 相似度阈值。
            verbose (bool): 是否打印详细信息 (当前未使用)。
        Returns:
            tuple: (filtered_texts, filtered_sims, all_sims)
                - filtered_texts (list[str]): 相似度 >= threshold 的文本。
                - filtered_sims (list[float]): 对应的相似度。
                - all_sims (list[float]): 所有文本的相似度。
        """
        if not texts: # Handle empty input
             return [], [], []

        similarities = self.calculate_similarity(image, texts) # [len(texts)]

        # Pair texts with similarities and filter
        filtered_pairs = [
            (text, sim) for text, sim in zip(texts, similarities)
            if sim >= threshold
        ]
        # Sort filtered pairs by similarity descending
        filtered_pairs.sort(key=lambda x: x[1], reverse=True)

        filtered_texts = [t[0] for t in filtered_pairs]
        filtered_sims = [t[1] for t in filtered_pairs]
        # Convert all similarities to list for consistency
        all_sims = similarities.tolist()

        return filtered_texts, filtered_sims, all_sims
